#  -*- coding: utf-8 -*-
import pandas as pd
from pymongo import UpdateOne,ASCENDING, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes,get_all_indexes_date,calc_negative_diff_dates,multi_computer
from util.database import DB_CONN
import time
from datetime import datetime, timedelta

"""
实现成交量均线的因子计算和保存
"""


class NewHighFactor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='new_high')

    def calc_hist_max_ratio(self,df,date):

        #print(date)
        max_value = df.loc[:date]['close'].max()
        max_ratio = round(df.loc[date]['close']/max_value,2)
        #print(date,max_value)
        return max_ratio


    def computer_single(self,code,is_index, autype, begin_date,end_date):
        dm = DataModule()
        start_time = time.time()
        print('计算新高比例, %s' % code)

        if begin_date is None:
            begin_date = "1990-12-19"

        df_daily = dm.get_k_data(code, index=is_index, autype=autype,begin_date = None,end_date= end_date)
        if df_daily.index.size > 0:
            df_daily.set_index(['date'], 1, inplace=True)
            update_requests = []

            df_trading_daily = df_daily.loc[df_daily.is_trading == True, :]
            df_trading_daily_copy = df_trading_daily.copy()
            try:
                # 计算阶段性最大值
                df_trading_daily_copy['ma5_max'] = df_trading_daily_copy['close'].rolling(5).max()
                df_trading_daily_copy['ma10_max'] = df_trading_daily_copy['close'].rolling(10).max()
                df_trading_daily_copy['ma20_max'] = df_trading_daily_copy['close'].rolling(20).max()
                df_trading_daily_copy['ma60_max'] = df_trading_daily_copy['close'].rolling(60).max()
                df_trading_daily_copy['ma99_max'] = df_trading_daily_copy['close'].rolling(99).max()
                df_trading_daily_copy['ma250_max'] = df_trading_daily_copy['close'].rolling(250).max()

                #计算当前价格相对于阶段性最大值的比例
                df_trading_daily_copy['ma5_max_ratio'] = round(df_trading_daily_copy['close']/df_trading_daily_copy['ma5_max'],2)
                df_trading_daily_copy['ma10_max_ratio'] = round(df_trading_daily_copy['close']/df_trading_daily_copy['ma10_max'],2)
                df_trading_daily_copy['ma20_max_ratio'] = round(df_trading_daily_copy['close']/df_trading_daily_copy['ma20_max'],2)
                df_trading_daily_copy['ma60_max_ratio'] = round(df_trading_daily_copy['close']/df_trading_daily_copy['ma60_max'],2)
                df_trading_daily_copy['ma99_max_ratio'] = round(df_trading_daily_copy['close']/df_trading_daily_copy['ma99_max'],2)
                df_trading_daily_copy['ma250_max_ratio'] = round(df_trading_daily_copy['close']/df_trading_daily_copy['ma250_max'],2)
                #历史最大值比例，要以截止计算日期前的数据进行计算
                #print(df_trading_daily_copy)
                df_trading_daily_copy['history_max_ratio'] = df_trading_daily_copy.apply(
                    lambda row: self.calc_hist_max_ratio(df_trading_daily_copy, row.name), axis=1)

                df_target_daily = df_trading_daily_copy.loc[begin_date:end_date]

                #df_target_daily.to_csv("newhigh_1122.csv")

                for date in df_target_daily.index:
                    update_requests.append(
                        UpdateOne(
                            {'code': code, 'date': date,'index': is_index},
                            {'$set': {'code': code,
                                      'date': date,
                                      'index': is_index,
                                      'ma5_max_ratio': df_target_daily.loc[date]['ma5_max_ratio'],
                                      'ma10_max_ratio': df_target_daily.loc[date]['ma10_max_ratio'],
                                      'ma20_max_ratio': df_target_daily.loc[date]['ma20_max_ratio'],
                                      'ma60_max_ratio': df_target_daily.loc[date]['ma60_max_ratio'],
                                      'ma99_max_ratio': df_target_daily.loc[date]['ma99_max_ratio'],
                                      'ma250_max_ratio': df_target_daily.loc[date]['ma250_max_ratio'],
                                      'history_max_ratio': df_target_daily.loc[date]['history_max_ratio']
                                      }},
                            upsert=True))
            except Exception as e:
                print(f'填充新高比例时发生错误，股票代码：{code} 是否指数：{is_index}，错误原因：{e}' , flush=True)

            if len(update_requests) > 0:
                update_result = self.collection.bulk_write(update_requests, ordered=False)
                end_time = time.time()

                print('填充新高比例，股票：%s，是否指数：%s,插入：%4d条，更新：%4d条,耗时：%.3f 秒' %
                      (code, is_index, update_result.upserted_count, update_result.modified_count, (end_time - start_time)),
                      flush=True)

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值(新高比例)，并保存到数据库中
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        self.collection.create_index([('code', 1), ('date', 1),('index',1)])

        #获取所有股票
        codes = get_all_codes()
        args = (False,'qfq',begin_date,end_date)
        multi_computer(computer_codes,codes,args)

        #获取所有指数
        codes = get_all_indexes_date()
        args = (True,None,begin_date,end_date)
        multi_computer(computer_codes,codes,args)


def computer_codes(codes,is_index,autype,begin_date,end_date):
    for code in codes:
        NewHighFactor().computer_single(code,is_index,autype,begin_date,end_date)


if __name__ == '__main__':
    # 执行因子的提取任务
    #hfq =HfqMAFactor()
    pd.set_option('display.width',500)
    pd.set_option('display.max_columns', 500)
    pd.set_option('display.max_colwidth', 500)
    NewHighFactor().compute('2020-11-21', '2020-11-24')
